from yandex_music import Client
from yandex_music.exceptions import NetworkError


class YandexMusicClient:
    def __init__(self, token):
        # Store the Yandex Music API token
        self.token = token
        # Initialize the Yandex Music client
        self.client = self.initialize_client()

    def initialize_client(self):
        # Initialize the Yandex Music client and handle network errors during initialization
        while True:
            try:
                return Client(token=self.token).init()
            except NetworkError:
                continue

    def get_playlists(self):
        # Prepare a list with a default playlist for "Liked" tracks
        playlists = [
            {
                "title": "Понравившиеся",
                "kind": -1,
                "owner_id": self.client.me.account.uid
            }
        ]

        # Retrieve custom playlists and add them to the list
        custom_playlists = self.client.users_playlists_list()
        for playlist_ in custom_playlists:
            playlist = {
                "title": playlist_.title,
                "kind": playlist_.kind,
                "owner_id": playlist_.owner.uid
            }
            playlists.append(playlist)

        return playlists

    def get_tracks_from_playlist(self, playlist):
        # Extract playlist details
        kind = playlist["kind"]
        owner_id = playlist["owner_id"]

        if kind == -1:
            # Get "Liked" tracks for the user
            short_tracks = self.client.users_likes_tracks()
            short_track_ids = [short_track.id for short_track in short_tracks]
            tracks = self.client.tracks(short_track_ids)
        else:
            # Get tracks from a custom playlist
            full_playlist = self.client.playlists_list(f"{owner_id}:{kind}")[0]
            short_tracks = full_playlist.fetch_tracks()
            tracks = [short_track.track for short_track in short_tracks]

        # Filter out unavailable tracks
        tracks = [track for track in tracks if track.available]

        return tracks

    def get_track_url(self, track_id):
        # Get the download URL for a track
        while True:
            try:
                track_info = self.client.tracks_download_info(track_id, True)
                # Sort tracks by bitrate and get the highest quality download link
                sorted_tracks = sorted(track_info, key=lambda x: x["bitrate_in_kbps"], reverse=True)
                track_url = sorted_tracks[0].get_direct_link()
                return track_url
            except NetworkError:
                continue
